import { MessageCreateParams } from "@anthropic-ai/sdk/resources"

const MAX_CACHING_MESSAGES = 4

// Caching strategy type - transforms a clean Anthropic body by adding cache_control
type CachingStrategy = (anthropicBody: MessageCreateParams) => MessageCreateParams

// Utility function to estimate token count
const estimateTokenCount = (text: string): number => Math.ceil(text.length / 4)

// Strategy 1: No Caching
const noCachingStrategy: CachingStrategy = (body) => body

// Strategy 2: System Messages Only
const systemOnlyStrategy: CachingStrategy = (body) => {
	let availableCacheMessages = MAX_CACHING_MESSAGES
	if (body.system && Array.isArray(body.system)) {
		return {
			...body,
			system: body.system.map((item) => {
				if (availableCacheMessages > 0) {
					availableCacheMessages -= 1
					return {
						...item,
						cache_control: { type: "ephemeral" },
					}
				}
				return item
			}),
		}
	}
	return body
}

// Strategy 3: System and Tools (High Impact)
const systemAndToolsStrategy: CachingStrategy = (body) => {
	const result = { ...body }
	let availableCacheMessages = MAX_CACHING_MESSAGES

	// Cache system messages
	if (result.system && Array.isArray(result.system)) {
		result.system = result.system.map((item) => {
			if (availableCacheMessages > 0) {
				availableCacheMessages -= 1
				return {
					...item,
					cache_control: { type: "ephemeral" },
				}
			}
			return item
		})
	}

	// Cache tool definitions
	if (result.tools && Array.isArray(result.tools) && result.tools.length > 0) {
		result.tools = result.tools.map((tool, index: number) => {
			if (index === result.tools!.length - 1 && availableCacheMessages > 0) {
				availableCacheMessages -= 1
				return {
					...tool,
					cache_control: { type: "ephemeral" },
				}
			}
			return tool
		})
	}

	return result
}

// Strategy 4: Optimized (Intelligent Caching)
const optimizedStrategy: CachingStrategy = (body) => {
	const result = { ...body }
	let availableCacheMessages = MAX_CACHING_MESSAGES

	// Always cache system messages
	if (result.system && Array.isArray(result.system)) {
		result.system = result.system.map((item) => {
			if (availableCacheMessages > 0) {
				availableCacheMessages -= 1
				return {
					...item,
					cache_control: { type: "ephemeral" },
				}
			}
			return item
		})
	}

	// Cache tool definitions
	if (result.tools && Array.isArray(result.tools) && result.tools.length > 0) {
		result.tools = result.tools.map((tool, index: number) => {
			if (index === result.tools!.length - 1 && availableCacheMessages > 0) {
				availableCacheMessages -= 1
				return {
					...tool,
					cache_control: { type: "ephemeral" },
				}
			}
			return tool
		})
	}

	// Cache large messages (>500 tokens)
	if (result.messages && Array.isArray(result.messages)) {
		result.messages = result.messages.map((message) => {
			if (message.content && typeof message.content === "string") {
				const tokens = estimateTokenCount(message.content)
				if (tokens > 500 && availableCacheMessages > 0) {
					availableCacheMessages -= 1
					return {
						...message,
						content: [
							{
								type: "text",
								text: message.content,
								cache_control: { type: "ephemeral" as const },
							},
						],
					}
				}
			} else if (message.content && Array.isArray(message.content)) {
				// Only add one cache control per message with array content
				let addedCacheControl = false
				const updatedContent = message.content.map((item) => {
					if (item.type === "text" && item.text) {
						const tokens = estimateTokenCount(item.text)
						if (tokens > 500 && availableCacheMessages > 0 && !addedCacheControl) {
							availableCacheMessages -= 1
							addedCacheControl = true
							return {
								...item,
								cache_control: { type: "ephemeral" as const },
							}
						}
					}
					return item
				})

				return {
					...message,
					content: updatedContent,
				}
			}
			return message
		})
	}

	return result
}

// Available caching strategies
export const CACHING_STRATEGIES = {
	none: noCachingStrategy,
	systemOnly: systemOnlyStrategy,
	systemAndTools: systemAndToolsStrategy,
	optimized: optimizedStrategy,
} as const

export type CachingStrategyName = keyof typeof CACHING_STRATEGIES
